/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */

package cn.com.sina.erp.easymr;

import com.hadoop.compression.lzo.DistributedLzoIndexer;
import com.hadoop.mapreduce.LzoTextInputFormat;
import java.io.IOException;
import java.net.URI;
import java.text.DateFormat;
import java.text.SimpleDateFormat;
import java.util.ArrayList;
import java.util.Calendar;
import java.util.Date;
import java.util.HashMap;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.CommandLineParser;
import org.apache.commons.cli.GnuParser;
import org.apache.commons.cli.HelpFormatter;
import org.apache.commons.cli.OptionBuilder;
import org.apache.commons.cli.Options;
import org.apache.commons.cli.ParseException;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.filecache.DistributedCache;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.mapreduce.Job;
import org.apache.hadoop.mapreduce.JobContext;
import org.apache.hadoop.mapreduce.Reducer;
import org.apache.hadoop.mapreduce.lib.input.TextInputFormat;
import org.apache.hadoop.mapreduce.lib.output.FileOutputFormat;
import org.apache.hadoop.util.Tool;
import org.apache.hadoop.util.ToolRunner;

/**
 *
 * @author guti
 * @version 0.2.0
 */
public class EasyMr {
    
    // argument parser
    public Options options = new Options();
    public CommandLine commandLine;
    private CommandLineParser commandLineParser = new GnuParser();
    // MapReduce objects
    public Configuration conf;
    public Job job;
    // behaviours
    public String mapOutputKeyType = "text"; // ["text", "textpair"]
    public Boolean compressMapperOutput = true;
    // user registered classes 
    public Map<String, Class> classMap = new HashMap<String, Class>();

    public String version = "0.2.0";
    public ArrayList<String> helpHeader = new ArrayList<String>();
    public ArrayList<String> helpFooter = new ArrayList<String>();
    
    
    public static String getProtectedArg(String key, JobContext context) {
        String protectedKey = "EasyMr_" + key;
        Configuration conf = context.getConfiguration();
        return conf.get(protectedKey);
    }

    public static String getCacheFilePath(String alias, JobContext context) throws IOException {
        String localPath = null;
        Configuration conf = context.getConfiguration();
        String filename = conf.get("EasyMr__Cache_" + alias);
        if (filename == null) {
            return null;
        }

        Path[] paths = DistributedCache.getLocalCacheFiles(conf);
        for (Path path : paths) {
            String[] fragments = path.toString().split("/");
            if (fragments[fragments.length - 1].equals(filename)) {
                localPath = path.toString();
                break;
            }
        }
        return localPath;
    }


    public EasyMr() {
        options.addOption(OptionBuilder.withLongOpt("help")
                .withDescription("Help message.").create());
        options.addOption(OptionBuilder.withLongOpt("mr").hasArg()
                .withDescription("MapReduce class name.").create());
        options.addOption(OptionBuilder.withLongOpt("input").hasOptionalArgs()
                .withDescription("Input path(s). Format: path,MapperClass[,lzo]").create());
        options.addOption(OptionBuilder.withLongOpt("output").hasArg()
                .withDescription("Output path.").create());
        options.addOption(OptionBuilder.withLongOpt("force")
                .withDescription("Force deletion of existing output path.").create());
        options.addOption(OptionBuilder.withLongOpt("cache").hasOptionalArgs()
                .withDescription("Cached file(s). Format: localpath,alias").create());
        options.addOption(OptionBuilder.withLongOpt("protectedargs").hasOptionalArgs()
                .withDescription("Protected arguments. Format: argkey=argvalue").create());
        options.addOption(OptionBuilder.withLongOpt("reducetasks").hasArg()
                .withDescription("Number of reduce tasks. Integer.").create());
        options.addOption(OptionBuilder.withLongOpt("jobname").hasArg()
                .withDescription("Job name.").create());
        options.addOption(OptionBuilder.withLongOpt("compress")
                .withDescription("Whether or not compress the output.").create());

        // deprecated
        options.addOption(OptionBuilder.withLongOpt("numreduces").hasArg().create());

        helpHeader.add("This is a Hadoop Mapreduce application powered by EasyMr.\n");
        helpHeader.add("The version of EasyMr is " + version + "\n");
        helpHeader.add("Important changes:\n");
        helpHeader.add(" * Option 'numreduces' is deprecated, use 'reducetasks' in replace.\n");
        helpHeader.add("\n");
        helpHeader.add("General arguments supported:\n");

        conf = new Configuration();
    }
    public void takeArgs(String[] args) throws ParseException, Exception {
        commandLine = commandLineParser.parse(options, args);
    }
    
    public void registerClass(Class cls) {
        String[] frags = cls.getName().split("[\\.\\$]");
        print("Registering class: " + frags[frags.length - 1] + " ...");
        classMap.put(frags[frags.length - 1], cls);
    }

    public int run(Class mrCls) throws Exception {
        if (commandLine.hasOption("help")) {
            String header = "";
            for (String line : helpHeader) {
                header += ">>  " + line;
            }

            String footer = ">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>\n";
            for (String line : helpFooter) {
                if (line.startsWith("Author") || line.startsWith("author")) {
                    footer += ">\n## " + line + " ##\n";
                    footer += "<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<\n";
                } else {
                    footer += "   " + line + " \\\n";
                }
            }

            HelpFormatter formatter = new HelpFormatter();
            formatter.printHelp("See below.", header, options,
                    "\nUsage:\n" + footer, false);
            return 0;
        }
        
        // conf and job
        if (compressMapperOutput) {
            print("Map output will get compressed.");
            conf.set("mapred.compress.map.output","true");
            conf.set("mapred.map.output.compression.codec","com.hadoop.compression.lzo.LzoCodec");
        }
        if (commandLine.hasOption("compress")) {
            print("Reduce output will get compressed.");
            conf.set("mapred.output.compress","true");
            conf.set("mapred.output.compression.codec","com.hadoop.compression.lzo.LzopCodec");
        }

        // setCacheFiles
        if (commandLine.hasOption("cache")) {
            for (String cache : commandLine.getOptionValues("cache")) {
                String[] cacheFrags = cache.split(",");
                // path
                URI path = new URI(cacheFrags[0]);
                // alias
                String alias = cacheFrags[1];
                String[] pathFrags = cacheFrags[0].split("/");
                String filename = pathFrags[pathFrags.length - 1];

                print("Adding cache: ");
                print("    path: " + cacheFrags[0]);
                print("    alias: " + alias);
                DistributedCache.addCacheFile(path, conf);
                conf.set("EasyMr__Cache_" + alias, filename);
            }
        }

        if (commandLine.hasOption("protectedargs")) {
            print("Setting protected args:");
            for (String protectedarg : commandLine.getOptionValues("protectedargs")) {
                int delimiterPosition = protectedarg.indexOf('=');
                String protectedargKey = protectedarg.substring(0, delimiterPosition);
                String protectedargValue = protectedarg.substring(delimiterPosition + 1);
                print("    key: " + protectedargKey);
                print("    value: " + protectedargValue);
                conf.set("EasyMr_" + protectedargKey, protectedargValue);
            }
        }

        // modifications taken upon conf will not affect job from here on
        job = new Job(conf);
        
        // setJarByClass
        print("Setting jar by class: " + mrCls.getName() + " ...");
        job.setJarByClass(mrCls);

        // setJobName
        String[] frags = mrCls.getName().split("[\\.\\$]");
        String jobname = frags[frags.length - 1];
        if (commandLine.hasOption("jobname")) {
            jobname += " " + commandLine.getOptionValue("jobname");
        }
        print("Setting job name: " + jobname + " ...");
        job.setJobName(jobname);

        // setMapOutputKeyClass
        print("Setting map output key type: " + mapOutputKeyType + " ...");
        if (mapOutputKeyType.equals("text")) {
            setTextAsMapperOutput();
        } else if (mapOutputKeyType.equals("textpair")) {
            setTextPairAsMapperOuput();
        }

        // setInputPaths
        if (commandLine.hasOption("input")) {
            for (String input : commandLine.getOptionValues("input")) {
                String[] inputFrags = input.split(",");

                // path
                String[] paths = parsePath(inputFrags[0]);
                // mapperClass
                Class mapperClass = classMap.get(inputFrags[1]);
                // [inputFormatClass]
                Class inputFormatClass = TextInputFormat.class;
                if (inputFrags.length == 3 && "lzo".equals(inputFrags[2])) {
                    inputFormatClass = LzoTextInputFormat.class;
                }

                for (String path : paths) {
                    Path p = new Path(path);
                    print("Adding input: ");
                    print("    path: " + p.toString());
                    print("    mapper: " + mapperClass.getName());
                    print("    format: " + inputFormatClass.getName());
                    MultipleInputs.addInputPath(job, p, inputFormatClass, mapperClass);
                }
            }
        }

        // setOutputPath
        if (commandLine.hasOption("output")) {
            String outputPath = commandLine.getOptionValue("output");
            print("Setting output path: " + outputPath + " ...");
            FileOutputFormat.setOutputPath(job, new Path(outputPath));

            if (commandLine.hasOption("force")) {
                FileSystem fs = FileSystem.get(job.getConfiguration());
                print("Deleting output path: " + outputPath + " ...");
                fs.delete(new Path(outputPath), true);
            }
        }
        

        // setReducerClass
        for (Class cls : classMap.values()) {
            if (Reducer.class.isAssignableFrom(cls)) {
                print("Setting reducer class: " + cls.getName() + " ...");
                job.setReducerClass(cls);
            }
        }

        // setNumReduceTasks
        int numReduceTasks = 1;
        if (commandLine.hasOption("reducetasks")) {
            numReduceTasks = Integer.parseInt(commandLine.getOptionValue("reducetasks"));
        } else if (commandLine.hasOption("numreduces")) {
            numReduceTasks = Integer.parseInt(commandLine.getOptionValue("numreduces"));
        }
        print("Setting number of reduce tasks: " + numReduceTasks + " ...");
        job.setNumReduceTasks(numReduceTasks);



        // job.waitForCompletion
        if (commandLine.hasOption("input")) {
            job.waitForCompletion(true);
        }

        // create lzo indexes
        if (commandLine.hasOption("compress")) {
            print("Indexing lzo compressed output ...");
            ToolRunner.run(new DistributedLzoIndexer(), new String[]{
                commandLine.getOptionValue("output")});
        }
        
        return 0;
    }

    public int run(String[] args) throws Exception {
        if (!commandLine.hasOption("mr")) {
            print("Option '--mr' is required. Available:");
            for (String cls : classMap.keySet()) {
                print (" * " + cls);
            }
            return 0;
        }
        String mrName = commandLine.getOptionValue("mr");
        if (classMap.containsKey(mrName)) {
            print("Delegating to mr class: " + mrName + " ...");
            ToolRunner.run((Tool) classMap.get(mrName).newInstance(), args);
        } else {
            print("mr class not found: " + mrName);
        }
        return 0;
    }

    private void print(String str) {
        if (str.endsWith("\n")) {
            System.out.print("{EasyMr} " + str);
        } else {
            System.out.println("{EasyMr} " + str);
        }
    }
    
    
    // method cluster for mapper output format
    private void setTextAsMapperOutput() {
        job.setMapOutputKeyClass(Text.class);
        job.setMapOutputValueClass(Text.class);
    }
    private void setTextPairAsMapperOuput() {
        job.setMapOutputKeyClass(TextPair.class);
        job.setMapOutputValueClass(Text.class);
        job.setPartitionerClass(TextPair.KeyPartitioner.class);
        job.setGroupingComparatorClass(TextPair.FirstComparator.class);
    }

    private String[] parsePath(String path) throws ParseException, java.text.ParseException {
        Pattern p0 = Pattern.compile("(\\{\\d{4}/\\d{2}/\\d{2}-\\d{4}/\\d{2}/\\d{2}})");
        Pattern p1 = Pattern.compile("\\d{4}/\\d{2}/\\d{2}");

        Matcher m0 = p0.matcher(path);
        if (m0.find()) {
            ArrayList<String> pathList = new ArrayList<String>();

            String pathPrefix = path.substring(0, m0.start());
            String pathSuffix = path.substring(m0.end());

            String dateRangeExp = m0.group();
            Matcher m1 = p1.matcher(dateRangeExp);
            m1.find();
            String dateStringBegin = m1.group();
            m1.find();
            String dateStringEnd = m1.group();

            DateFormat format = new SimpleDateFormat("yyyy/MM/dd");

            Date dateBegin = (Date)format.parse(dateStringBegin);
            Calendar calendarBegin = Calendar.getInstance();
            calendarBegin.setTime(dateBegin);
            Date dateEnd = (Date)format.parse(dateStringEnd);
            Calendar calendarEnd = Calendar.getInstance();
            calendarEnd.setTime(dateEnd);

            if (calendarEnd.before(calendarBegin)) {
                calendarBegin.setTime(dateEnd);
                calendarEnd.setTime(dateBegin);
            }
            while (!calendarEnd.before(calendarBegin)) {
                pathList.add(pathPrefix + format.format(calendarBegin.getTime()) + pathSuffix);
                calendarBegin.add(Calendar.DATE, 1);
            }

            String[] paths = new String[pathList.size()];
            return pathList.toArray(paths);
        } else {
            return new String[]{path};
        }
    }
    
}
